[HIP] [JIT] fp8_mqa_logits: hand-written gfx950 prefill indexer kernel - #5046

Validation status: PR kernel commit d046a14c3b, rebased onto ROCm/aiter@9aa8a6b9.
The original op-level sweep below was measured before that rebase. The separately reported GLM-5.2-MXFP4 E2E and accuracy results are now included below under GLM-5.2-MXFP4 E2E and accuracy (2026-09-05 update), using vLLM 3ff4f02df and AITER 0.1.21.post1 plus this PR's files.

Motivation

The DeepSeek-V3.2 / GLM-5 "lightning indexer" produces the sparse-attention
selection logits. For each query row m and KV position n:

logits[m, n] = sum_h relu(<Q[m, h, :], K[n, :]>) * kv_scale[n] * weights[m, h]
               for n in [cu_seqlen_ks[m], cu_seqlen_ke[m]), -inf elsewhere

In the prefill path the caller has already gathered K out of the paged cache into
a contiguous [N, 128] buffer with a separate [N] fp32 scale, so there is no
block table and no preshuffle. This PR adds a hand-written HIP kernel for that
path on gfx950 (CDNA4), alongside the existing Triton/Gluon
aiter.ops.triton.attention.fp8_mqa_logits. The paged decode half is #5047.

cu_seqlen_ks is a request's base offset into the gathered K buffer and
cu_seqlen_ke grows by one per query row, so each request is a causal triangle
and several requests arrive packed into one [M, N] chunk. Chunks are sized so
M*N*4 stays under VLLM_SPARSE_INDEXER_MAX_LOGITS_MB, which makes a 32K
request (4096, 32768) and a 128K one (1024, 131072).

Technical Details

New op: aiter/ops/fp8_mqa_logits.py

fp8_mqa_logits is a drop-in for the Triton entry point -- same tensors
(q_fp8, k_fp8, kv_scale, weights, cu_seqlen_ks, cu_seqlen_ke) and the
same clean_logits semantics. Key design points:

• No LDS staging. K is streamed HBM->register and contracted 32 heads x 32
  tokens per mfma_scale_f32_32x32x64_f8f6f4 tile.
• BLOCK_M query rows share one K stream, so K is read once per row block
  rather than once per row. Grid is (ceil(M / BLOCK_M), SplitN); SplitN
  splits each block's KV tile range so a small row grid still fills the device.
• Warps split the tile range, not the rows. That keeps several independent K
  streams in flight per block. Splitting rows across warps instead -- so K is read
  once per BLOCK_M * NUM_WARPS rows rather than once per BLOCK_M -- measured
  1.5-23% slower: the latency hiding is worth more than the L2 traffic it saves.
• Paired-row v_permlane32_swap_b32 head reduce. v_permlane32_swap_b32 a, b
  leaves a's reduction in lanes 0-31 and b's in lanes 32-63, so two query rows
  reduce with one swap and one add, and the store that follows uses all 64 lanes
  instead of half.
• kv-scale hoist. kv_scale >= 0 and ReLU is positive-homogeneous, so the
  scale is applied once per KV column after the head reduction rather than inside
  it.
• Row groups dispatched high-m first for N > 2048. Under causal masking a
  row group's work grows with m, so in natural order the longest-running blocks
  are dispatched LAST and become the tail; reversing starts them first and lets
  the short ones fill in behind. Worth +6% at (4096, 4096) and +3% at
  (8192, 8192), neutral elsewhere, and free -- it only reorders block dispatch.
• Fused -inf fill. With clean_logits the kernel writes the -inf outside
  each row's window itself. The Triton path pre-fills all M*N elements with
  torch.full and then overwrites the valid ones, paying for the valid region
  twice.
• -fno-honor-nans for the module, so the ReLU is a single v_max_f32.
  Without it LLVM must assume a signalling NaN and emits an IEEE canonicalize
  first -- two VALU per accumulator value, which is ~27% of the kernel's VALU.

BlockM, SplitN, num_warps, unroll2 and reverse_rows are all tunable;
zero means "use the host heuristic".

The kernel is gfx950-only and fixed at n_heads=32, head_dim=128 -- the shipped
GLM-5-FP8 indexer shape. is_supported(num_heads, head_dim) gates on that, so a
caller that also serves other shapes can route them to the Triton kernel rather
than trip a TORCH_CHECK. This mirrors how _should_use_asm_kernel gates the
head_size=128-only ASM paged-attention kernel in aiter/ops/attention.py.

Files added / changed:

• aiter/ops/fp8_mqa_logits.py -- the op and its support gate
• csrc/kernels/fp8_mqa_logits.cu -- kernel and host dispatch
• csrc/include/fp8_mqa_logits.h, csrc/pybind/fp8_mqa_logits_pybind.cu
• csrc/include/rocm_ops.hpp, aiter/jit/optCompilerConfig.json -- module_fp8_mqa_logits
• op_tests/test_fp8_mqa_logits.py -- correctness + perf sweep

Test Plan

op_tests/test_fp8_mqa_logits.py runs Triton and HIP on identical inputs and
grades both against one fp32 torch reference -- the same ref_fp8_mqa_logits the
Triton lane's test uses. Gates are an exact -inf mask match plus
calc_diff < 1e-3 and checkAllclose; tolerances are not widened.

The sweep is the cartesian product of 14 (s_q, s_k) shapes,
num_heads in {32,64,128}, head_dim in {64,128}, clean_logits in {0,1} and
six window modes -- 900 cases, 150 of which the HIP kernel supports. Points worth
calling out:

• Six window modes. Beyond causal and cp, the sweep covers misaligned,
  empty (rows with cu_ends below zero or below cu_starts), past_end
  (bounds beyond seq_len_kv) and multi_req (several requests packed into one
  chunk, so cu_starts jumps at each boundary and a block's rows straddle it).
  All are legal indexer input at a chunk boundary, and all are where the masking
  and the -inf fill are easiest to get wrong. Grading the HIP kernel under
  past_end caught two out-of-bounds writes during development, both on the
  clean_logits=False path.
• NaN-poisoned output. The output buffer is filled with NaN before each call,
  so a position the kernel fails to write fails the -inf mask check. Without it
  the check is close to vacuous -- the caching allocator hands back a block a
  previous case already left holding the correct -inf.
• The real chunk shapes: (4096, 32768), (2048, 65536), (1024, 131072).
  The reference runs in query-row chunks so its [heads, s_q, s_k] score tensor
  stays bounded (unchunked it is 17 GiB at heads=32, s_q=1024, s_k=131072).
• Cases the HIP kernel does not support leave its columns nan rather than
  reporting a wrong-but-fast number, and any case dropped for lack of memory is
  logged by name so a short table cannot read as full coverage.

python3 op_tests/test_fp8_mqa_logits.py

Test Result

All correctness gates pass on gfx950 across the sweep; per-case hip err matches
triton err. Grading the HIP kernel under past_end caught two out-of-bounds
writes on the clean_logits=False path (a cu_starts past seq_len_kv left the
fill's first range unclamped, and the store bounded abs_pos only by cu_ends);
both are fixed here.

Performance on MI355x/gfx950, num_heads=32, head_dim=128, run_perftest on an
otherwise idle GPU. causal is one request per chunk, multi_req is four:

Summarised over the 32-shape sweep (both clean_logits settings, both windows):

The win tracks s_k / s_q, which is what the shared-K-stream design predicts:
BLOCK_M rows amortise one K read, so the longer a row's KV range is relative to
the row grid, the more there is to amortise. Every shape with s_k >= 8 * s_q is
a win (1.19x - 6.39x), and the largest is the small-M/long-KV corner
(128, 32768) at 6.4x, where the row grid alone cannot fill the device and
SplitN does the work.

The losses are the square, short-KV shapes -- (8192, 8192) at 0.97x/0.83x and
(1024, 1024)/(4096, 4096) under clean_logits=False -- where the tile loop is
short relative to the per-block Q/weights prologue. That prologue is also what
pins the kernel at 2 waves/SIMD: BLOCK_M=4 needs 194 VGPR, of which Q and the
per-row weights are 128, and both scale with BLOCK_M, so KV reuse and register
pressure cannot be traded apart in this design (BLOCK_M=8 needs 256 VGPR and
spills 113). Staging K through LDS would decouple them; until then
is_supported() plus a shape check lets a caller keep Triton on that corner.

clean_logits=True is the better case for the HIP kernel (1.64x vs 1.43x), which
is the fused -inf fill showing up: the Triton path pays a torch.full over all
s_q * s_k elements before the kernel overwrites the valid ones.

Related Work

• [FlyDSL] gfx950 FP8 MQA logits indexer kernel #4538 targets the same prefill operator with a FlyDSL implementation and broader shape coverage.
• feat(gfx950): config-gated BLOCK_Q fp8_mqa_logits for DSA indexer prefill #4180 targets the same operator with a Gluon BLOCK_Q implementation.

This PR proposes a structurally different, fixed-shape HIP backend that streams K directly from HBM to registers and shares the K stream across query rows. A direct head-to-head comparison with the active alternatives and quantitative roofline utilization remain pending.

Current Validation

• Rebased onto upstream main at 9aa8a6b91f1972952314bd16176a76d392f6b85c with no overlapping-file conflicts.
• black==26.3.0 --check: pass for the added Python op and test.
• ruff==0.16.0 check: pass for the added Python op and test.
• python3 -m py_compile: pass for the added Python op and test.
• A standalone op-suite rerun on the rebased head, direct competing-PR A/B, and quantitative roofline analysis remain pending. The supplied model-level integration report is included below.

AI assistance: OpenAI Codex was used for rebase adaptation, formatting, static checks, overlap research, and PR drafting. The submitter reviewed the resulting commit.

GLM-5.2-MXFP4 E2E and accuracy (2026-09-05 update)

The following supplied experiment report addresses the E2E and accuracy request in #5046 (comment). These are reported measurements on the integration stack specified below; this documentation update did not rerun the experiments.

E2E and accuracy for this PR on GLM-5.2-MXFP4 / MI355X (gfx950) / TP4, per @nholmber's request.

Headline: 1.07x output throughput at 128k input, at every concurrency point, with accuracy held.
No measurable change at 1k input, which is expected — this is a prefill kernel and 1k prefill is a
small share of E2E work.

Setup

Integration — one env-gated branch in rocm_fp8_mqa_logits()
(vllm/v1/attention/ops/rocm_aiter_mla_sparse.py), since this kernel is a drop-in for the
aiter.ops.triton.attention.fp8_mqa_logits that vLLM already calls there. GLM-5.2 is
index_n_heads=32, index_head_dim=128, so is_supported() holds and the kernel really engages.

Both arms ran concurrently on one node, one 4-GPU half each, so they share wall-clock and
thermals; the only difference is the kernel:

The branch raises rather than silently falling back, and the server log confirms which kernel
each arm used: 4 provenance lines on the "after" arm (one per TP rank), 0 on "before". So the
"after" column cannot be an accidental re-measurement of "before".

128k in / 1k out

TTFT 1.07-1.13x better, TPOT 1.05-1.07x better.

Independent confirmation from wall-clock (not derived from the same JSON metrics) — each
vllm bench serve invocation, before -> after:

1k in / 1k out

Per-run wall-clock is identical to within a second at 5 of 7 points.

TTFT is omitted for this scenario: at conc 16-64 the "after" arm's points happened to run while the
other arm was loading its 408 GB of weights, and that host-side NVMe/CPU burst shows up in TTFT
(0.78-0.85x) while leaving throughput and TPOT at 1.00x. The clean points either side (conc 4, 8,
128, 256) are all 1.00x on TTFT too, so we read the dip as contention, not kernel. Happy to re-run
the 1k sweep synchronised if the TTFT column matters.

Accuracy

Validation gates and expectations from the AMD "GLM-5.2 MXFP4 - vLLM status and validation" deck.

GPQA differs by one question out of 198 (0.5 pp), well inside the +-2.1 pp standard error at that
sample size. RULER is at the ceiling on both arms at both lengths, i.e. the sparse path still selects
correctly with the new kernel.

Notes

• Why 1.07x E2E when the op-level geomean is 1.53-1.79x: the indexer prefill logits kernel is one
  component of E2E time, so its speedup is diluted by everything else in the step. A ~6-7% E2E gain at
  long context from a single kernel is the expected shape of this result, and we'd suggest not setting
  E2E expectations from the op-level numbers.
• Why a prefill kernel moves TPOT by 5-7%: with chunked prefill, prefill chunks are interleaved
  with decode batches, so faster prefill means decode steps queue behind it for less time. This effect
  only shows up E2E.
• Nominal vs actual concurrency at 128k: weights are ~102 GB/GPU at TP4, leaving ~157 GB for KV at
  default utilisation, and this model costs ~55 KB/token — about 21 requests of 132k. So conc 32 and 64
  queue rather than running fully in parallel, which is why throughput saturates at ~105-112 tok/s from
  conc 16 on. Identical on both arms, so the comparison holds, but the concurrency label is nominal.
• Prefix caching disabled (--no-enable-prefix-caching). With it on, a repeated prompt skips
  prefill entirely, which would hide exactly what this PR changes. This does mean absolute numbers are
  not directly comparable to runs that leave it at the default.
• Not compared against [FlyDSL] gfx950 FP8 MQA logits indexer kernel #4538 or [Triton/Gluon] [gfx950] add an optimized prefill fp8_mqa_logits for H64D128 #5048. @nholmber also asked for those; this run was scoped to
  before/after against the reference AITER kernel. We can add [FlyDSL] gfx950 FP8 MQA logits indexer kernel #4538 (the FlyDSL implementation of the
  same kernel) as a third arm if it would help the decision.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.